/*
 * Copyright 2022 The Open Islands Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.openislands.oi.worker.manager;

import com.baomidou.mybatisplus.core.toolkit.support.SFunction;
import com.fasterxml.jackson.core.type.TypeReference;
import org.openislands.oi.config.MyNodeInfo;
import org.openislands.oi.constant.*;
import org.openislands.oi.dao.entity.JobDO;
import org.openislands.oi.dao.entity.TaskDO;
import org.openislands.oi.error.LogicException;
import org.openislands.oi.manager.CommonJobManager;
import org.openislands.oi.manager.CommonOfNodeTaskManager;
import org.openislands.oi.manager.CommonOperatorManager;
import org.openislands.oi.pojo.dto.JobDTO;
import org.openislands.oi.pojo.dto.OfNodeTaskDTO;
import org.openislands.oi.pojo.dto.OperatorDTO;
import org.openislands.oi.pojo.dto.TaskDTO;
import org.openislands.oi.scheduling.DAGOperator;
import org.openislands.oi.scheduling.StateEngine;
import org.openislands.oi.util.ObjectUtils;
import org.openislands.oi.util.ValueCopyUtils;
import org.openislands.oi.worker.scheduler.OperatorRuntimeManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.dao.DuplicateKeyException;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import javax.annotation.Resource;
import java.util.Date;
import java.util.List;
import java.util.Map;

@Service
public class TaskManager {
    private static final Logger log = LoggerFactory.getLogger(TaskManager.class);

    @Resource
    private CommonJobManager commonJobManager;

    @Resource
    private CommonOperatorManager commonOperatorManager;

    @Resource
    private CommonOfNodeTaskManager commonOfNodeTaskManager;

    @Resource
    private OperatorRuntimeManager operatorRuntimeManager;

    @Resource
    private StateEngine stateEngine;
    @Resource
    private MyNodeInfo myNodeInfo;

    public Boolean createJob(JobDTO jobDTO) {
        try {
            this.fillResourceParams(jobDTO);
            this.saveConfs(jobDTO);
            final JobDO jobDO = ValueCopyUtils.copyProperties(jobDTO, new JobDO());
            jobDO.setProgress(0);
            jobDO.setElapsed(0);
            jobDO.setCreateTime(new Date());
            return commonJobManager.getJobManager().save(jobDO);
        } catch (DuplicateKeyException e) {
            return true;
        }
    }

    private void fillResourceParams(JobDTO jobDTO) {
        //todo: engine address
    }

    private void saveConfs(JobDTO jobDTO){
        //todo: save conf/conf overwrite/runtime conf
    }

    public JobDTO queryJob(String jobId) {
        return commonJobManager.queryJob(jobId);
    }

    public List<JobDTO> queryJobByState(JobState jobState, String leaderNode) {
        JobDTO jobDTO = new JobDTO();
        jobDTO.setState(jobState);
        jobDTO.setLeaderNode(leaderNode);
        return this.queryAnyJob(jobDTO);
    }

    public List<JobDTO> queryJobByState(JobState jobState) {
        JobDTO jobDTO = new JobDTO();
        jobDTO.setState(jobState);
        return this.queryAnyJob(jobDTO);
    }

    public List<JobDTO> queryAnyJob(JobDTO jobDTO) {
        return commonJobManager.getJobManager().queryAnyLikeDTO(jobDTO);
    }

    public Boolean updateJob(String jobId, JobDTO jobDTO) {
        jobDTO.setJobId(jobId);
        return this.updateAnyJob(jobDTO, JobDO::getJobId);
    }

    @SafeVarargs
    public final <Res> Boolean updateAnyJob(JobDTO jobDTO, SFunction<JobDO, Res>... functions) {
        return commonJobManager.updateAnyJob(jobDTO, functions);
    }

    public MessageCodeEnum updateJobState(String jobId, JobState toState) {
        return commonJobManager.updateJobState(jobId, toState);
    }

    public Boolean deleteJob(String jobId) {
        return commonJobManager.deleteJob(jobId);
    }

    public Boolean stopJob(String jobId, JobState toState) {
        OfNodeTaskDTO ofNodeTaskDTO = new OfNodeTaskDTO();
        ofNodeTaskDTO.setJobId(jobId);
        //todo: get latest
        List<OfNodeTaskDTO> ofNodeTaskDTOS = commonOfNodeTaskManager.queryAnyOfNodeTask(ofNodeTaskDTO);
        if (ofNodeTaskDTOS.isEmpty()) {
            throw new LogicException(MessageCodeEnum.DATA_NOT_FOUND_ERROR, "can not found job");
        }
        TaskState toTaskState = TaskState.valueOf(toState.name());
        boolean stopSt = true;
        for (OfNodeTaskDTO ofNodeTask : ofNodeTaskDTOS) {
            stopSt = stopSt && this.stopTask(jobId, ofNodeTask.getTaskName(), ofNodeTask.getTaskVersion(), ofNodeTask.getRole(), toTaskState);
        }
        return stopSt;
    }

    @Transactional
    public Boolean createTask(TaskDTO taskDTO) {
        final TaskDO taskDO = ValueCopyUtils.copyProperties(taskDTO, new TaskDO());
        try {
            taskDO.setProgress(0);
            taskDO.setElapsed(0);
            taskDO.setCreateTime(new Date());
            if (!commonJobManager.getTaskManager().save(taskDO)) {
                throw new LogicException(MessageCodeEnum.DATA_INSERT_ERROR, "create task at node failed");
            }
        } catch (DuplicateKeyException e) {
            log.warn("insert task success, but duplicate");
        }
        //todo: Operators used by the node should be pre-coordinated, but not here
        OperatorDTO operatorDTO = commonOperatorManager.getOperatorForTask(taskDTO);
        if (operatorDTO == null) {
            throw new LogicException(MessageCodeEnum.EXECUTION_ERROR, "can not found operator");
        }
        final OfNodeTaskDTO ofNodeTaskDTO = ValueCopyUtils.copyProperties(taskDTO, new OfNodeTaskDTO());
        ofNodeTaskDTO.setOperator(operatorDTO);
        switch (ofNodeTaskDTO.getOperator().getPackagedType()) {
            case PYTHON:
            case JAR:
                ofNodeTaskDTO.setContainerType(OperatorContainerType.WORKER_SUBPROCESS);
                break;
            case IMAGE:
                ofNodeTaskDTO.setContainerType(OperatorContainerType.PROCESS_IN_DOCKER);
                break;
            default:
                throw new LogicException(MessageCodeEnum.EXECUTION_ERROR, String.format("not support %s packaged type", ofNodeTaskDTO.getOperator().getPackagedType()));
        }
        //todo
        ofNodeTaskDTO.setCores(4);
        ofNodeTaskDTO.setMemory("4g");
        ofNodeTaskDTO.setNodeId(myNodeInfo.getNodeId());
        Map<RoleType, List<Map<String, DAGOperator>>> taskConfig = ObjectUtils.json2Object(taskDO.getTaskConfig(), new TypeReference<Map<RoleType, List<Map<String, DAGOperator>>>>() {
        });
        assert taskConfig != null;
        taskConfig.forEach(((role, nodes) -> {
            nodes.forEach(nodeConfigs -> {
                if (nodeConfigs.containsKey(myNodeInfo.getNodeId())) {
                    ofNodeTaskDTO.setRole(role);
                    ofNodeTaskDTO.setTaskConfig(ObjectUtils.object2Json(nodeConfigs.get(myNodeInfo.getNodeId())));
                    try {
                        boolean st = commonOfNodeTaskManager.createOfNodeTask(ofNodeTaskDTO);
                        if (!st) {
                            throw new LogicException(MessageCodeEnum.DATA_INSERT_ERROR, String.format("create task at node(%s) role(%s) failed", myNodeInfo.getNodeId(), role));
                        }
                    } catch (DuplicateKeyException e) {
                        log.warn("insert task at node success, but duplicate");
                    }
                }

            });
        }));
        return true;
    }

    public Boolean stopTask(String jobId, String taskName, Integer taskVersion, RoleType role, TaskState toState) {
        List<OfNodeTaskDTO> ofNodeTasks = commonOfNodeTaskManager.queryOfNodeTask(jobId, taskName, taskVersion, role, myNodeInfo.getNodeId());
        if (ofNodeTasks.isEmpty()) {
            throw new LogicException(MessageCodeEnum.DATA_NOT_FOUND_ERROR, "can not found task at node");
        } else if (ofNodeTasks.size() > 1) {
            throw new LogicException(MessageCodeEnum.DATA_REDUNDANCY_ERROR, "can not more than one task at the same time");
        }
        OfNodeTaskDTO ofNodeTask = ofNodeTasks.get(0);
        boolean terminateSt = operatorRuntimeManager.terminateRuntime(ofNodeTask);
        if (terminateSt) {
            MessageCodeEnum updateStateSt = commonOfNodeTaskManager.updateState(jobId, taskName, taskVersion, role, myNodeInfo.getNodeId(), toState);
            if (updateStateSt == MessageCodeEnum.SUCCESS || updateStateSt == MessageCodeEnum.STATE_UPDATE_NOT_EFFECTIVE || updateStateSt == MessageCodeEnum.STATE_TRANSITION_FAILED) {
                ofNodeTask.setEndTime(new Date());
                boolean updateOfNodeTaskSt = commonOfNodeTaskManager.updateOfNodeTask(jobId, taskName, taskVersion, role, myNodeInfo.getNodeId(), ofNodeTask);
                return true;
            } else {
                throw new LogicException(MessageCodeEnum.STATE_UPDATE_ERROR, "terminate runtime success, but update state failed");
            }
        } else {
            throw new LogicException(MessageCodeEnum.COORDINATION_COMMAND_EXECUTION_ERROR, "terminate runtime failed");
        }
    }

    public Boolean deleteTaskByJob(String jobId) {
        return commonJobManager.deleteTaskByJob(jobId);
    }

    public Boolean deleteTask(String jobId, String taskName, Integer taskVersion) {
        return commonJobManager.deleteTask(jobId, taskName, taskVersion);
    }

    public TaskDTO queryTask(String jobId, String taskName, Integer taskVersion) {
        TaskDTO taskDTO = new TaskDTO();
        taskDTO.setJobId(jobId);
        taskDTO.setTaskName(taskName);
        taskDTO.setTaskVersion(taskVersion);
        List<TaskDTO> taskDTOList = this.queryAnyTask(taskDTO);
        if (taskDTOList.isEmpty()) {
            return null;
        } else if (taskDTOList.size() > 1) {
            throw new LogicException(MessageCodeEnum.DATA_REDUNDANCY_ERROR, "find more than one");
        }
        return taskDTOList.get(0);
    }

    public List<TaskDTO> queryTaskByJob(String jobId) {
        TaskDTO taskDTO = new TaskDTO();
        taskDTO.setJobId(jobId);
        return this.queryAnyTask(taskDTO);
    }

    public List<TaskDTO> queryTaskByState(TaskState taskState) {
        TaskDTO taskDTO = new TaskDTO();
        taskDTO.setState(taskState);
        return this.queryAnyTask(taskDTO);
    }

    public List<TaskDTO> queryAnyTask(TaskDTO taskDTO) {
        return commonJobManager.queryAnyTask(taskDTO);
    }

    public Boolean updateTask(String jobId, String taskName, Integer taskVersion, TaskDTO taskDTO) {
        taskDTO.setJobId(jobId);
        taskDTO.setTaskName(taskName);
        taskDTO.setTaskVersion(taskVersion);
        return this.updateAnyTask(taskDTO, TaskDO::getJobId, TaskDO::getTaskName, TaskDO::getTaskVersion);
    }

    public MessageCodeEnum updateTaskState(String jobId, String taskName, Integer taskVersion, TaskState toState) {
        return commonJobManager.updateTaskState(jobId, taskName, taskVersion, toState);
    }

    @SafeVarargs
    public final <Res> Boolean updateAnyTask(TaskDTO taskDTO, SFunction<TaskDO, Res>... functions) {
        return commonJobManager.updateAnyTask(taskDTO, functions);
    }
}
